#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import struct

from ambari_ws4py.framing import (
  Frame,
  OPCODE_CONTINUATION,
  OPCODE_TEXT,
  OPCODE_BINARY,
  OPCODE_CLOSE,
  OPCODE_PING,
  OPCODE_PONG,
)
from ambari_ws4py.compat import unicode, py3k

__all__ = [
  "Message",
  "TextMessage",
  "BinaryMessage",
  "CloseControlMessage",
  "PingControlMessage",
  "PongControlMessage",
]


class Message(object):
  def __init__(self, opcode, data=b"", encoding="utf-8"):
    """
    A message is a application level entity. It's usually built
    from one or many frames. The protocol defines several kind
    of messages which are grouped into two sets:

    * data messages which can be text or binary typed
    * control messages which provide a mechanism to perform
      in-band control communication between peers

    The ``opcode`` indicates the message type and ``data`` is
    the possible message payload.

    The payload is held internally as a a :func:`bytearray` as they are
    faster than pure strings for append operations.

    Unicode data will be encoded using the provided ``encoding``.
    """
    self.opcode = opcode
    self._completed = False
    self.encoding = encoding

    if isinstance(data, unicode):
      if not encoding:
        raise TypeError("unicode data without an encoding")
      data = data.encode(encoding)
    elif isinstance(data, bytearray):
      data = bytes(data)
    elif not isinstance(data, bytes):
      raise TypeError(f"{type(data)} is not a supported data type")

    self.data = data

  def single(self, mask=False):
    """
    Returns a frame bytes with the fin bit set and a random mask.

    If ``mask`` is set, automatically mask the frame
    using a generated 4-byte token.
    """
    mask = os.urandom(4) if mask else None
    return Frame(body=self.data, opcode=self.opcode, masking_key=mask, fin=1).build()

  def fragment(self, first=False, last=False, mask=False):
    """
    Returns a :class:`ws4py.framing.Frame` bytes.

    The behavior depends on the given flags:

    * ``first``: the frame uses ``self.opcode`` else a continuation opcode
    * ``last``: the frame has its ``fin`` bit set
    * ``mask``: the frame is masked using a automatically generated 4-byte token
    """
    fin = 1 if last is True else 0
    opcode = self.opcode if first is True else OPCODE_CONTINUATION
    mask = os.urandom(4) if mask else None
    return Frame(body=self.data, opcode=opcode, masking_key=mask, fin=fin).build()

  @property
  def completed(self):
    """
    Indicates the the message is complete, meaning
    the frame's ``fin`` bit was set.
    """
    return self._completed

  @completed.setter
  def completed(self, state):
    """
    Sets the state for this message. Usually
    set by the stream's parser.
    """
    self._completed = state

  def extend(self, data):
    """
    Add more ``data`` to the message.
    """
    if isinstance(data, bytes):
      self.data += data
    elif isinstance(data, bytearray):
      self.data += bytes(data)
    elif isinstance(data, unicode):
      self.data += data.encode(self.encoding)
    else:
      raise TypeError(f"{type(data)} is not a supported data type")

  def __len__(self):
    return len(self.__unicode__())

  def __str__(self):
    if py3k:
      return self.data.decode(self.encoding)
    return self.data

  def __unicode__(self):
    return self.data.decode(self.encoding)


class TextMessage(Message):
  def __init__(self, text=None):
    Message.__init__(self, OPCODE_TEXT, text)

  @property
  def is_binary(self):
    return False

  @property
  def is_text(self):
    return True


class BinaryMessage(Message):
  def __init__(self, bytes=None):
    Message.__init__(self, OPCODE_BINARY, bytes, encoding=None)

  @property
  def is_binary(self):
    return True

  @property
  def is_text(self):
    return False

  def __len__(self):
    return len(self.data)


class CloseControlMessage(Message):
  def __init__(self, code=1000, reason=""):
    data = b""
    if code:
      data += struct.pack("!H", code)
    if reason is not None:
      if isinstance(reason, unicode):
        reason = reason.encode("utf-8")
      data += reason

    Message.__init__(self, OPCODE_CLOSE, data, "utf-8")
    self.code = code
    self.reason = reason

  def __str__(self):
    if py3k:
      return self.reason.decode("utf-8")
    return self.reason

  def __unicode__(self):
    return self.reason.decode(self.encoding)


class PingControlMessage(Message):
  def __init__(self, data=None):
    Message.__init__(self, OPCODE_PING, data)


class PongControlMessage(Message):
  def __init__(self, data):
    Message.__init__(self, OPCODE_PONG, data)
